from tqdm import tqdm
import torch
import torch.nn.functional as F


class BaseTrainer:
    def __init__(self, model, optimizer, scheduler, device='cpu', clip_grad=1.0, patience=5):
        self.model = model.to(device)
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.device = device
        self.clip_grad = clip_grad  # 梯度裁剪阈值
        self.patience = patience  # 早停容忍轮数
        self.best_val_acc = 0.0  # 最佳验证准确率
        self.early_stop_counter = 0  # 早停计数器

    def train_epoch(self, dataloader):
        # 保持原训练逻辑不变
        self.model.train()
        total_loss = 0

        for data, target in dataloader:
            data, target = data.to(self.device), target.to(self.device)

            # 前向传播
            outputs = self.model(data)
            loss = F.cross_entropy(outputs, target)

            # 反向传播
            self.optimizer.zero_grad()
            loss.backward()

            # 梯度裁剪
            if self.clip_grad is not None:
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    max_norm=self.clip_grad
                )

            # 参数更新
            self.optimizer.step()
            total_loss += loss.item()

        # 学习率调度
        self.scheduler.step()
        return total_loss / len(dataloader)

    def validate(self, dataloader):
        """新增验证方法：计算准确率并检查早停"""
        self.model.eval()
        total_correct = 0
        total_samples = 0

        with torch.no_grad():
            for data, target in dataloader:
                data, target = data.to(self.device), target.to(self.device)
                outputs = self.model(data)
                _, predicted = torch.max(outputs.data, 1)  # 取预测概率最大的类别
                total_samples += target.size(0)
                total_correct += (predicted == target).sum().item()

        # 计算准确率
        val_acc = total_correct / total_samples

        # 早停逻辑：若当前准确率未超过最佳值，计数器+1；否则更新最佳值并重置计数器
        if val_acc <= self.best_val_acc:
            self.early_stop_counter += 1
            early_stop = self.early_stop_counter >= self.patience
        else:
            self.best_val_acc = val_acc
            self.early_stop_counter = 0
            early_stop = False

        return val_acc, early_stop